import torch
import os
import torch.nn as nn 
import torch.nn.functional as F
import onnxruntime as ort
import numpy as np


def write_data(data, filepath):
    data = data.reshape(-1)
    with open(filepath, 'w') as f:
        for i in range(len(data)):
            f.write('{}\n'.format(data[i]))


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.relu = nn.ReLU()
        self.pooling = nn.AvgPool2d(kernel_size=2)
        self.flatten = nn.Flatten()
        self.conv = nn.Conv2d(kernel_size=3, in_channels=3, out_channels=3, padding=1)
        ninputs = 3 * 4 * 4
        noutputs = 10
        self.linear = nn.Linear(ninputs, noutputs)

    def forward(self, x):
        out = x
        out = self.conv(out)
        out = self.pooling(out)
        out = self.flatten(out)
        out = self.linear(out)
        out = self.relu(out)
        return out


class SampleOnnx(object):
    def __init__(self, onnxpath):
        # provider = ['CPUExecutionProvider']
        self.session = ort.InferenceSession(onnxpath)
    
    def infer(self, x):
        input_dict = {
            'input': x,
        }
        output_names = ['output']
        outputs = self.session.run(output_names, input_dict)
        return outputs[0]


def export_onnx(model, onnx_path):
    inputs = torch.randn(1, 3, 8, 8)
    torch.onnx.export(model, inputs, onnx_path, 
        verbose=True,
        opset_version=11, 
        input_names=['input'], 
        output_names=['output'], 
    )


def test_onnx():
    onnxpath = '../data/sample.onnx'
    # if not os.path.exists(onnxpath):
    if True:
        model = Model()
        export_onnx(model, onnxpath)
    inputs = np.ones((1, 3, 8, 8), dtype=np.float32)
    demo = SampleOnnx(onnxpath)
    y = demo.infer(inputs)
    print(y)
    print(y.shape)
    print(np.sum(y))


if __name__ == '__main__':
    test_onnx()
